from pathlib import Path
import re
from enum import Enum
from typing import Optional
from matplotlib import pyplot as plt
from .tview import tview_bam, BamTview


class PointOption(Enum):
    A = {'c': '#9DE093', 'marker': '$A$'}
    T = {'c': '#DF8D8F', 'marker': '$T$'}
    G = {'c': '#CC99CC', 'marker': '$G$'}
    C = {'c': '#7EA2ED', 'marker': '$C$'}
    D = {'c': '#808080', 'marker': '$-$'}


class VariantPointOption(Enum):
    A = {'c': '#00AD43', 'marker': '$A$'}
    T = {'c': '#CE2029', 'marker': '$T$'}
    G = {'c': "#660099", 'marker': '$G$'}
    C = {'c': "#0000FF", 'marker': '$C$'}
    D = {'c': '#808080', 'marker': '$-$'}


def get_point_option(base: str, is_var: bool) -> dict:
    option = VariantPointOption if is_var else PointOption
    if base == 'A':
        return option.A.value
    if base == 'T':
        return option.T.value
    if base == 'G':
        return option.G.value
    if base == 'C':
        return option.C.value
    if base == '-':
        return option.D.value
    return {'c': '#000000', 'marker': f'${base}$'}


def set_base_xy(seq: list[str], y: int, x_dict: dict, y_dict: dict, var_x_dict: dict, var_y_dict: dict, ref: str):
    for i, base in enumerate(seq):
        if base != " ":
            if ref[i] == base:
                x_dict.setdefault(base, list()).append(i + 1)
                y_dict.setdefault(base, list()).append(y)
            else:
                var_x_dict.setdefault(base, list()).append(i + 1)
                var_y_dict.setdefault(base, list()).append(y)


def plot_snv(tview: BamTview, extend: int, img: str | Path, dpi=300):
    length, depth = len(tview.reference), len(tview.reads)
    plt.figure(figsize=(length / 8, depth / 6))
    plt.xlim((0, length + 1))
    plt.ylim((-depth - 1), 4)
    plt.xticks([])
    yticks = [2, 1] + list(range(-10, -depth - 1, -10))
    ylabels = ['Reference', 'Consensus'] + [f'{abs(i)}X' for i in range(-10, -depth - 1, -10)]
    plt.yticks(ticks=yticks, labels=ylabels, size=7.5)
    plt.hlines(y=0, xmin=0, xmax=length+1, color='#808080')
    ax = plt.gca()
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["bottom"].set_visible(False)
    x_dict, y_dict, var_x_dict, var_y_dict = dict(), dict(), dict(), dict()
    set_base_xy(tview.reference, 2, x_dict, y_dict, var_x_dict, var_y_dict, tview.reference)
    set_base_xy(tview.consensus, 1, x_dict, y_dict, var_x_dict, var_y_dict, tview.reference)
    for i, read in enumerate(tview.reads):
        set_base_xy(read, -(i+1), x_dict, y_dict, var_x_dict, var_y_dict, tview.reference)
    for base in x_dict.keys():
        plt.scatter(x_dict[base], y_dict[base], **get_point_option(base, False))
    for base in var_x_dict.keys():
        plt.scatter(var_x_dict[base], var_y_dict[base], **get_point_option(base, True))
    plt.arrow(extend+tview.reference[:extend].count(' ')+1, 4, 0, -1, head_width=0.2, lw=2, color='red', length_includes_head=True)
    plt.savefig(img, dpi=dpi, bbox_inches='tight')
